from environment.state import State
from network.sub_network import LSTMNet, CNNLSTMNet, DuelingNoisyLSTMNet
from network.memory import Memory
import torch
import torch.nn as nn
import numpy as np


class PERDQN(object):
    def __init__(self, n_state_feature, n_neurons, n_actions, memory_capacity, lr, epsilon, env_a_shape,
                 target_replace_iter, batch_size, gamma, n_layer=1):
        self.n_feature = n_state_feature
        self.memory_capacity = memory_capacity

        self.eval_net, self.target_net = LSTMNet(self.n_feature[-1], n_neurons, n_actions, n_layer=n_layer), \
                                         LSTMNet(self.n_feature[-1], n_neurons, n_actions, n_layer=n_layer)
        # self.eval_net, self.target_net = CNNLSTMNet(n_length=self.n_feature[0], n_feature=self.n_feature[1],
        #                                             n_neurons=n_neurons, n_actions=n_actions), \
        #                                  CNNLSTMNet(n_length=self.n_feature[0], n_feature=self.n_feature[1],
        #                                             n_neurons=n_neurons, n_actions=n_actions)
        # self.eval_net, self.target_net = DuelingNoisyLSTMNet(self.n_feature[-1], n_neurons, n_actions, n_layer=n_layer),\
        #                                  DuelingNoisyLSTMNet(self.n_feature[-1], n_neurons, n_actions, n_layer=n_layer)
        self.learn_step_counter = 0  # for target updating
        self.memory_counter = 0  # for storing memory
        self.memory = Memory(self.memory_capacity)
        self.optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=lr)
        self.epsilon = epsilon
        self.env_a_shape = env_a_shape
        self.n_actions = n_actions
        self.target_replace_iter = target_replace_iter
        self.batch_size = batch_size
        self.gamma = gamma

        np.random.seed(0)

    def choose_action(self, x):
        x = torch.unsqueeze(torch.FloatTensor(x), 0)
        # input only one sample
        if np.random.uniform() < self.epsilon:  # greedy
            actions_value = self.eval_net.forward(x)
            action = torch.max(actions_value, 1)[1].data.numpy()
            action = action[0] if self.env_a_shape == 0 else action.reshape(self.env_a_shape)  # return the argmax index
        else:  # random
            action = np.random.randint(0, self.n_actions)
            action = action if self.env_a_shape == 0 else action.reshape(self.env_a_shape)
        print("action: ", action)
        return action

    def store_transition(self, s: State, a, r, s_: State, done):
        info = {
            's': s.feature, 'a': a, 'r': r, 's_': s_.feature, "finished": done
        }
        # replace the old memory with new memory
        self.memory.store(info)
        self.memory_counter += 1

    def learn(self):
        # target parameter update
        if self.learn_step_counter % self.target_replace_iter == 0:
            self.target_net.load_state_dict(self.eval_net.state_dict())
        self.learn_step_counter += 1

        # sample batch transitions
        b_s, b_a, b_r, b_s_, dones, tree_idx, is_weights = self.get_sample_from_memory()
        b_a = torch.reshape(b_a, (-1, 1))
        b_r = torch.reshape(b_r, (-1, 1))
        dones = torch.reshape(dones, (-1, 1))
        # q_eval w.r.t the action in experience
        q_values = self.eval_net(b_s).gather(1, b_a)  # shape (batch, 1)
        next_q_values = self.eval_net(b_s_)
        q_next = self.target_net(b_s_).detach()  # detach from graph, don't backpropagate

        ### regular q learning
        # q_target = b_r + self.gamma * torch.max(q_next, 1).values.view(self.batch_size, 1)  # shape (batch, 1)

        ### double q learning
        next_q_value = q_next.gather(1, torch.max(next_q_values, 1)[1].unsqueeze(1))
        q_target = b_r + self.gamma * next_q_value * (1 - dones)

        loss = self.loss_func(q_values, q_target, is_weights)

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        # self.eval_net.reset_noise()
        # self.target_net.reset_noise()

        abs_errors = torch.abs(q_target - q_values).detach().numpy()
        self.memory.batch_update(tree_idx, abs_errors)

    def get_sample_from_memory(self):
        tree_idx, minibatch, ISWeights = self.memory.sample(self.batch_size)
        b_s = []
        b_a = []
        b_r = []
        b_s_ = []
        dones = []
        for dic in minibatch:
            b_s.append(dic['s'])
            b_a.append(dic['a'])
            b_r.append(dic['r'])
            b_s_.append(dic['s_'])
            dones.append(dic["finished"])
        b_s = np.array(b_s)
        b_a = np.array(b_a)
        b_r = np.array(b_r)
        b_s_ = np.array(b_s_)
        dones = np.array(dones)
        return torch.FloatTensor(b_s), torch.LongTensor(b_a), torch.FloatTensor(b_r), torch.FloatTensor(b_s_), \
               torch.FloatTensor(dones), tree_idx, ISWeights

    def save_model(self, path, name=None):
        if name is None:
            name = "eval_net.pt"
        torch.save(self.eval_net.state_dict(), path + "/" + name)

    def load_model(self, path, name=None):
        if name is None:
            name = "eval_net.pt"
        self.eval_net.load_state_dict(torch.load(path + "/" + name))

    @staticmethod
    def loss_func(values, target, weights):
        diff = torch.FloatTensor(values - target)
        return torch.mean(diff * diff * torch.tensor(weights))


if __name__ == '__main__':
    print(np.random.choice(3, 10))
